"""
Hi-C 扩散模型可视化工具
支持：
1. 逐步加噪过程可视化（数据预处理后 → 逐步加噪到纯噪声）
2. 逐步去噪过程可视化（纯噪声 → 逐步去噪到重建样本）
3. 训练/验证样本对比可视化
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import torch
import matplotlib
from matplotlib import rcParams
from .hic_datasets import calculate_insulation_score

# 尝试启用中文字体，防止中文显示为方框
rcParams['font.sans-serif'] = [
    'Noto Sans CJK SC', 'SimHei', 'WenQuanYi Zen Hei',
    'Source Han Sans SC', 'Microsoft YaHei', 'Arial Unicode MS',
    'DejaVu Sans'
]
rcParams['font.family'] = 'sans-serif'
rcParams['axes.unicode_minus'] = False


def visualize_forward_diffusion(
    x0,
    betas,
    num_timesteps=16,
    timestep_strategy="linspace",
    timestep_stride=0,
    max_timestep=None,
    use_mean=False,
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=None,
    titles=None,
    suptitle=None,
    figsize_per_col=3.5,
    dpi=150,
):
    """
    可视化前向扩散过程 q(x_t | x_0)
    
    Args:
        x0: 原始数据 [H, W] 或 [B, H, W]，若多样本则展示第一个
        betas: beta schedule，numpy array [T]
        num_timesteps: 要展示的时间步数量（包含 t=0）
        timestep_strategy: 'linspace', 'random', 'manual'
        timestep_stride: >0 时按步长选择 t
        max_timestep: 限制最大 t
        use_mean: True=只展示均值，False=加随机噪声
        log_display: True=log1p 显示
        signed_display: True=支持正负值（RdBu_r），False=非负（Reds）
        per_panel_scale: True=每子图独立色阶，False=全局统一
        save_path: 保存路径，None 则显示
        titles: 自定义每个子图标题
        suptitle: 总标题
        figsize_per_col: 每列宽度
        dpi: 分辨率
        
    Returns:
        fig: matplotlib figure
    """
    # 处理输入形状
    if isinstance(x0, torch.Tensor):
        x0 = x0.detach().cpu().numpy()
    if x0.ndim == 3:
        x0 = x0[0]  # 取第一个样本
    assert x0.ndim == 2, f"x0 must be 2D, got shape {x0.shape}"
    
    # 计算 alpha_bar
    alphas = 1.0 - betas
    alphas_cumprod = np.cumprod(alphas)
    T = len(betas)
    
    # 前向加噪函数
    def q_sample(x0_np, t_idx):
        at = float(alphas_cumprod[int(t_idx)])
        mean = np.sqrt(at) * x0_np
        if use_mean:
            return mean
        std = np.sqrt(1.0 - at)
        noise = np.random.randn(*x0_np.shape).astype(np.float32)
        return mean + std * noise
    
    # 选择时间步
    t_list = _select_timesteps(
        T, num_timesteps, timestep_strategy, timestep_stride, max_timestep
    )
    
    # 生成各时间步的加噪结果
    matrices = [x0]
    default_titles = ["t=0 (x₀)"]
    for t in t_list:
        xt = q_sample(x0, int(t))
        matrices.append(xt)
        alpha_t = alphas_cumprod[int(t)]
        default_titles.append(f"t={int(t)}\n(ᾱ={alpha_t:.3f})")
    
    if titles is None:
        titles = default_titles
    
    if suptitle is None:
        suptitle = "前向扩散过程：x₀ → 纯噪声"
    
    # 绘制网格
    fig = _plot_matrix_grid(
        matrices, 
        titles, 
        suptitle=suptitle,
        log_display=log_display,
        signed_display=signed_display,
        per_panel_scale=per_panel_scale,
        figsize_per_col=figsize_per_col,
        dpi=dpi,
    )
    
    if save_path:
        os.makedirs(os.path.dirname(save_path) or '.', exist_ok=True)
        fig.savefig(save_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)
    
    return fig


def visualize_reverse_diffusion(
    model,
    diffusion,
    shape,
    device='cuda',
    num_timesteps=16,
    timestep_strategy="linspace",
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=None,
    model_kwargs=None,
    suptitle=None,
    figsize_per_col=3.5,
    dpi=150,
):
    """
    可视化反向去噪过程 p(x_{t-1} | x_t)
    
    Args:
        model: 去噪模型
        diffusion: 扩散过程对象
        shape: 样本形状 [B, C, H, W]
        device: 设备
        num_timesteps: 要展示的时间步数量
        其他参数同 visualize_forward_diffusion
        
    Returns:
        fig: matplotlib figure
    """
    model.eval()
    
    # 选择要记录的时间步
    T = diffusion.num_timesteps
    t_list = _select_timesteps(T, num_timesteps, timestep_strategy, 0, None)
    t_set = set(t_list)
    
    # 逐步去噪并记录
    matrices = []
    titles = []
    
    with torch.no_grad():
        img = torch.randn(*shape, device=device)
        matrices.append(img[0, 0].cpu().numpy())
        titles.append(f"t={T} (纯噪声)")
        
        for i in range(T)[::-1]:
            t = torch.tensor([i] * shape[0], device=device)
            out = diffusion.p_sample(
                model, img, t, 
                clip_denoised=True,
                model_kwargs=model_kwargs
            )
            img = out["sample"]
            
            # 记录选定的时间步
            if i in t_set or i == 0:
                matrices.append(img[0, 0].cpu().numpy())
                titles.append(f"t={i}")
    
    if suptitle is None:
        suptitle = "反向去噪过程：纯噪声 → x₀"
    
    # 绘制网格
    fig = _plot_matrix_grid(
        matrices, 
        titles, 
        suptitle=suptitle,
        log_display=log_display,
        signed_display=signed_display,
        per_panel_scale=per_panel_scale,
        figsize_per_col=figsize_per_col,
        dpi=dpi,
    )
    
    if save_path:
        os.makedirs(os.path.dirname(save_path) or '.', exist_ok=True)
        fig.savefig(save_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)
    
    model.train()
    return fig


def visualize_comparison(
    original, 
    generated, 
    log_display=True,
    save_path=None,
    figsize=(12, 5),
    dpi=150,
):
    """
    对比原始样本与生成样本
    
    Args:
        original: 原始 Hi-C 矩阵 [H, W]
        generated: 生成的 Hi-C 矩阵 [H, W]
        log_display: 是否 log 显示
        save_path: 保存路径
        
    Returns:
        fig: matplotlib figure
    """
    if isinstance(original, torch.Tensor):
        original = original.detach().cpu().numpy()
    if isinstance(generated, torch.Tensor):
        generated = generated.detach().cpu().numpy()
    
    # 准备显示数据
    if log_display:
        orig_disp = np.log1p(original)
        gen_disp = np.log1p(generated)
    else:
        orig_disp = original
        gen_disp = generated
    
    # 计算差异
    diff = generated - original
    
    # 统一色阶
    vmin = min(np.percentile(orig_disp, 2), np.percentile(gen_disp, 2))
    vmax = max(np.percentile(orig_disp, 98), np.percentile(gen_disp, 98))
    
    # 绘制
    fig, axes = plt.subplots(1, 3, figsize=figsize, dpi=dpi)
    
    im0 = axes[0].imshow(orig_disp, cmap='Reds', origin='lower', vmin=vmin, vmax=vmax)
    axes[0].set_title('原始 (Original)')
    axes[0].axis('off')
    plt.colorbar(im0, ax=axes[0], fraction=0.046)
    
    im1 = axes[1].imshow(gen_disp, cmap='Reds', origin='lower', vmin=vmin, vmax=vmax)
    axes[1].set_title('生成 (Generated)')
    axes[1].axis('off')
    plt.colorbar(im1, ax=axes[1], fraction=0.046)
    
    im2 = axes[2].imshow(diff, cmap='RdBu_r', origin='lower')
    axes[2].set_title('差异 (Difference)')
    axes[2].axis('off')
    plt.colorbar(im2, ax=axes[2], fraction=0.046)
    
    fig.suptitle('Hi-C 样本对比', fontsize=14, y=0.98)
    fig.tight_layout()
    
    if save_path:
        os.makedirs(os.path.dirname(save_path) or '.', exist_ok=True)
        fig.savefig(save_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)
    
    return fig


def _contact_decay_curve(matrix):
    """计算接触衰减曲线：按主对角线距离d统计均值。"""
    x = np.asarray(matrix, dtype=np.float64)
    n = x.shape[0]
    dists = np.arange(n)
    means = np.zeros(n, dtype=np.float64)
    for d in range(n):
        diag_vals = np.diagonal(x, offset=d)
        if diag_vals.size == 0:
            means[d] = np.nan
        else:
            # 忽略非正值以减少噪声影响
            pos = diag_vals[diag_vals > 0]
            if pos.size > 0:
                means[d] = pos.mean()
            else:
                means[d] = np.nanmean(diag_vals)
    # log1p 以压缩量级，便于可视化
    with np.errstate(invalid='ignore'):
        means = np.log1p(means)
    return dists, means


def visualize_comparison_with_metrics(
    original,
    generated,
    log_display=False,
    save_path=None,
    figsize=(20, 4),
    dpi=150,
    insulation_window=10,
    signed_display=True,
):
    """
    对比原始与生成样本，并附带：
      - 接触衰减 (contact decay)
      - 绝缘分数 (insulation score)
    图像布局：原始 | 生成 | 差异 | 衰减曲线 | 绝缘分数
    """
    if isinstance(original, torch.Tensor):
        original = original.detach().cpu().numpy()
    if isinstance(generated, torch.Tensor):
        generated = generated.detach().cpu().numpy()

    # 展示矩阵（支持符号对数）
    if log_display:
        if signed_display:
            to_disp = lambda m: np.sign(m) * np.log1p(np.abs(m))
        else:
            to_disp = lambda m: np.log1p(np.maximum(m, 0))
    else:
        to_disp = lambda m: m

    orig_disp = to_disp(original)
    gen_disp = to_disp(generated)
    diff = generated - original  # 差异图保留线性差值，便于正负偏差判断

    # 统一色阶
    if signed_display:
        cat_abs = np.abs(np.concatenate([orig_disp.flatten(), gen_disp.flatten()]))
        cat_abs = cat_abs[np.isfinite(cat_abs)]
        a = float(np.percentile(cat_abs, 98.0)) if cat_abs.size > 0 else 1.0
        a = max(a, 1e-6)
        vmin, vmax = -a, +a
        cmap_core = 'RdBu_r'
    else:
        vmin = min(np.percentile(orig_disp, 2), np.percentile(gen_disp, 2))
        vmax = max(np.percentile(orig_disp, 98), np.percentile(gen_disp, 98))
        cmap_core = 'Reds'

    # 差异色阶对称
    diff_disp = diff
    b = float(np.percentile(np.abs(diff_disp.flatten()), 98.0)) if diff_disp.size else 1.0
    b = max(b, 1e-6)
    vmin_diff, vmax_diff = -b, +b

    fig, axes = plt.subplots(1, 5, figsize=figsize, dpi=dpi)

    im0 = axes[0].imshow(orig_disp, cmap=cmap_core, origin='lower', vmin=vmin, vmax=vmax)
    axes[0].set_title('原始')
    axes[0].axis('off')
    plt.colorbar(im0, ax=axes[0], fraction=0.046)

    im1 = axes[1].imshow(gen_disp, cmap=cmap_core, origin='lower', vmin=vmin, vmax=vmax)
    axes[1].set_title('生成')
    axes[1].axis('off')
    plt.colorbar(im1, ax=axes[1], fraction=0.046)

    im2 = axes[2].imshow(diff_disp, cmap='RdBu_r', origin='lower', vmin=vmin_diff, vmax=vmax_diff)
    axes[2].set_title('差异')
    axes[2].axis('off')
    plt.colorbar(im2, ax=axes[2], fraction=0.046)

    # 接触衰减曲线
    d_o, m_o = _contact_decay_curve(original)
    d_g, m_g = _contact_decay_curve(generated)
    axes[3].plot(d_o, m_o, label='原始', color='#d62728')
    axes[3].plot(d_g, m_g, label='生成', color='#1f77b4')
    axes[3].set_title('接触衰减（log1p均值）')
    axes[3].set_xlabel('对角距离 (bins)')
    axes[3].set_ylabel('平均联系 (log1p)')
    axes[3].grid(True, alpha=0.3)
    axes[3].legend()

    # 绝缘分数
    try:
        ins_o = calculate_insulation_score(np.asarray(original, dtype=np.float64), window_size=insulation_window)
        ins_g = calculate_insulation_score(np.asarray(generated, dtype=np.float64), window_size=insulation_window)
    except Exception:
        # 兜底：简单滑窗替代，防止导入问题导致崩溃
        def _ins(mat, w=insulation_window):
            n = mat.shape[0]
            out = np.zeros(n)
            for i in range(w, n-w):
                up = mat[i-w:i, i-w:i].mean()
                down = mat[i:i+w, i:i+w].mean()
                cross = mat[i-w:i, i:i+w].mean()
                denom = 0.5 * (up + down) if (up + down) > 0 else 1.0
                out[i] = cross / denom
            return out
        ins_o = _ins(np.asarray(original, dtype=np.float64))
        ins_g = _ins(np.asarray(generated, dtype=np.float64))

    axes[4].plot(ins_o, label='原始', color='#d62728')
    axes[4].plot(ins_g, label='生成', color='#1f77b4')
    axes[4].set_title(f'绝缘分数 (window={insulation_window})')
    axes[4].set_xlabel('基因组位置 (bins)')
    axes[4].set_ylabel('分数')
    axes[4].grid(True, alpha=0.3)
    axes[4].legend()

    fig.suptitle('Hi-C 对比与指标', fontsize=14, y=0.98)
    fig.tight_layout()

    if save_path:
        os.makedirs(os.path.dirname(save_path) or '.', exist_ok=True)
        fig.savefig(save_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)

    return fig


def _select_timesteps(T, num_show, strategy, stride, max_t):
    """选择要展示的时间步"""
    num_show = max(1, min(num_show - 1, T))  # -1 because we add t=0 separately
    
    if stride and stride > 0:
        grid = np.arange(0, T, stride, dtype=int)
        t_list = grid[:num_show]
    elif strategy == 'random':
        t_list = np.random.choice(T, size=num_show, replace=False)
        t_list.sort()
    else:  # linspace
        t_list = np.linspace(0, T - 1, num_show, dtype=int)
    
    # 移除 t=0（因为会单独添加）
    t_list = t_list[t_list > 0]
    
    if max_t and max_t > 0:
        t_list = t_list[t_list <= min(max_t, T - 1)]
    
    return t_list


def _plot_matrix_grid(
    matrices, 
    titles, 
    suptitle="",
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    figsize_per_col=3.5,
    dpi=150,
):
    """绘制矩阵网格"""
    num_mats = len(matrices)
    cols = min(6, num_mats)
    rows = int(np.ceil(num_mats / cols))
    
    # 准备显示数据
    if signed_display:
        disp = []
        for m in matrices:
            d = np.sign(m) * np.log1p(np.abs(m)) if log_display else m
            disp.append(d)
        cmap = 'RdBu_r'
    else:
        disp = [np.log1p(m) if log_display else m for m in matrices]
        cmap = 'Reds'
    
    # 计算色阶
    if per_panel_scale:
        vmins = [float(np.percentile(d, 2.0)) for d in disp]
        vmaxs = [float(np.percentile(d, 98.0)) for d in disp]
    else:
        cat = np.concatenate([d.flatten() for d in disp])
        vmin_global = float(np.percentile(cat, 2.0))
        vmax_global = float(np.percentile(cat, 98.0))
    
    # 绘制
    fig, axes = plt.subplots(rows, cols, figsize=(figsize_per_col * cols, figsize_per_col * rows), dpi=dpi)
    if rows == 1 and cols == 1:
        axes = [axes]
    else:
        axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
    
    for i in range(rows * cols):
        ax = axes[i]
        if i < len(disp):
            if per_panel_scale:
                im = ax.imshow(disp[i], cmap=cmap, origin='lower', vmin=vmins[i], vmax=vmaxs[i])
            else:
                im = ax.imshow(disp[i], cmap=cmap, origin='lower', vmin=vmin_global, vmax=vmax_global)
            ax.set_title(titles[i], fontsize=10)
            ax.axis('off')
            # 添加色条
            plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        else:
            ax.axis('off')
    
    if suptitle:
        fig.suptitle(suptitle, fontsize=14, y=0.98)
    
    fig.tight_layout()
    return fig


def visualize_training_batch_forward(
    batch,
    betas,
    save_dir,
    step,
    num_samples=2,
    num_timesteps=16,
    **kwargs
):
    """
    训练循环中调用的批量可视化函数
    
    Args:
        batch: 数据批次，包含 'hic' 键
        betas: beta schedule
        save_dir: 保存目录
        step: 当前训练步数
        num_samples: 可视化样本数
        num_timesteps: 时间步数
        **kwargs: 传给 visualize_forward_diffusion 的其他参数
    """
    os.makedirs(save_dir, exist_ok=True)
    
    hic_data = batch['hic']  # [B, 1, H, W]
    if isinstance(hic_data, torch.Tensor):
        hic_data = hic_data.detach().cpu().numpy()
    
    num_samples = min(num_samples, hic_data.shape[0])
    
    for b in range(num_samples):
        x0 = hic_data[b, 0]  # [H, W]
        
        save_path = os.path.join(save_dir, f"forward_step_{step:06d}_sample_{b}.png")
        
        # 添加样本信息到标题
        cell_id = batch.get('cell_id', ['unknown'])[b] if 'cell_id' in batch else 'unknown'
        timepoint = batch.get('timepoint', ['unknown'])[b] if 'timepoint' in batch else 'unknown'
        suptitle = f"前向扩散 | Step {step} | Sample {b}\nCell: {cell_id} | Timepoint: {timepoint}"
        
        visualize_forward_diffusion(
            x0=x0,
            betas=betas,
            num_timesteps=num_timesteps,
            save_path=save_path,
            suptitle=suptitle,
            **kwargs
        )

